#include <stdio.h>

/* use these for random number generation */
#include <stdlib.h>
// #include <time.h>

#include "../../big_math/hdr/big_math.h"
#include "../hdr/rsa.h"


big_int* Ti_b;
big_int* temp;

big_int* Zi_b;
big_int* Ziplus_b;
big_int* Pi_b;
big_int* Piplus_b;

void rsa_init() {
   lib_init();

   Ti_b = (big_int*) malloc(sizeof(big_int));
   temp = (big_int*) malloc(sizeof(big_int));

   Zi_b = (big_int*) malloc(sizeof(big_int));
   Ziplus_b = (big_int*) malloc(sizeof(big_int));
   Pi_b = (big_int*) malloc(sizeof(big_int));
   Piplus_b = (big_int*) malloc(sizeof(big_int));
}



/* e is the input
   prime (p or q) is the output */
int generate_prime(const int rsa_e_value,
                   int *p_prime_array,
                   const int int_length_of_p) {
   int i = 0;
   for (i = 0; i < int_length_of_p; i++) {
      p_prime_array[i] = rand();
   }

   /* at this point we have a big random number...
      since this project is about optimizing the
      modular exponentiation in the RSA encryption
      algorithm, we won't bother checking for
      primality right now */
   
   /* after verifying the primality, we should
      also check that gcd(rsa_e_value, prime - 1)
      is 1, which would require our big math routines
      for sure */


}

int bit_length(int p) {
   int result = 0;
   int mask = 0;
   int i = sizeof(int) * 8 - 1;

   for (i; i >= 0; i--) {
      mask = 1 << i;

      if ((mask & p) != 0) {
         break;
      }
   }

   return i + 1;
}


/* calculates: cres = (ares * bres * rinv) % n */

int mont_prod(int ares, int bres, int np, int n, int r) {
   int Ti = 0;
   int Tiplus = 0;
   int mi = 0;
   int i = 0;
   int result = 0;
   int len_r = 0;

   len_r = bit_length(r);

   /* loop from 0 to l - 1, where 2^l = r */
   for (i = 0; i < len_r - 1; i++) {
      mi = ((Ti & 1) + ((ares >> i) & 1) * (bres & 1)) % 2;
      Tiplus = (Ti + ((ares >> i) & 1) * bres + mi * n) / 2;
      Ti = Tiplus;
   }
   if (Tiplus >= n) {
      result = Tiplus - n;
   }
   else {
      result = Tiplus;
   }

   return result;
}


// calculates montgomery product using big integers
// note: non-reentrant
// result: Tiplus_b
void mont_prod_big(big_int* Tiplus_b, big_int* ares, big_int* bres, big_int* n, big_int* r) {
   int mi;
   int i;
   int len_r;
   len_r = num_bits(r);

   for(i = 0; i < len_r - 1; i++) {
      mi = (bit_at(Ti_b, 0) + bit_at(ares, i) * bit_at(bres, 0) ) % 2;
      add(temp, bit_mult(n, mi), bit_mult(bres, bit_at(ares, i)));
      add(Tiplus_b, Ti_b, temp);
      assign (Ti_b, Tiplus_b);
   }

   if ( less_than ( n, Tiplus_b) ) {
      subtract(Tiplus_b, Tiplus_b, n);
   }
}





/* m is the message
   k is the exponent
   n is the modulus
   r*rinv + n*np = 1
   j is the bit-length of k
*/
int mod_exp(int m, int k, int np, int n, int r) {
   int mres = 0;
   int Pi = 0;  /* the modded square in loop i */
   int Piplus = 0;  /* the modded square for the next loop */
   int zres = 0; /* n-residue of the output in loop i */
   int ziplus = 0;  /* n-residue of the output in the next loop */
   int i = 0;
   int len_k = 0;

   int num_square_mont_prod = 0;
   int num_mult_mont_prod = 0;

   len_k = bit_length(k);
   printf("k = %x, len_k = %d\n", k, len_k);

   mres = (m*r)%n; /* immediately take the n-residue
                    * of the input message
		    */

   zres = r%n; /* now take the n-residue of the starting
                * output value, 1
		*/

   printf("mres: %d\n", mres);
   printf("zres: %d\n", zres);

   Pi = mres;

   for (i = 0; i < len_k; i++) {
      num_square_mont_prod++;
      Piplus = mont_prod(Pi, Pi, np, n, r);
      // printf("Piplus %da: %d\n", i, Piplus);
      if (((1 << i) & k) != 0) {
         num_mult_mont_prod++;
         ziplus = mont_prod(zres, Pi, np, n, r);
	 // printf("ziplus %db: %d\n", i, ziplus);
      }
      else {
         ziplus = zres;
      }
      /* Piplus is the new Pi,
       * ziplus is the new zres
       */
      Pi = Piplus;
      zres = ziplus;
   }

   printf("square mont prods: %d\n", num_square_mont_prod);
   printf("mult mont prods: %d\n", num_mult_mont_prod);
   return mont_prod(zres, 1, np, n, r);
}


big_int* mod_exp_big(big_int* mres, big_int* k, big_int*n, big_int* r, big_int* zres) {

   int len_k = num_bits(k);
   int i;

   assign(Pi_b, mres);
   assign(Zi_b, zres);

   for (i = 0; i < len_k; i++) {
      
      mont_prod_big(Piplus_b, Pi_b, Pi_b, n, r);
      // mont prod pi * pi
      if (bit_at(k, i) != 0) {
         mont_prod_big(Ziplus_b, Zi_b, Pi_b, n ,r);
         // mont prod zres * pi
      } else {
         assign(Ziplus_b, Zi_b);
         //ziplus = zres
      }

      assign (Pi_b, Piplus_b);
      assign (Zi_b, Ziplus_b);
   }

   return Zi_b;
}
